import torch
import argparse
from torch.utils.data import DataLoader
from torchvision import transforms
from dataset import Flowers
from model import Vit
from runner import Runner


def main(arg):
    device = torch.device(arg.device if torch.cuda.is_available() else "cpu")
    # 建立数据集
    data_transform = {
        "train": transforms.Compose([transforms.RandomResizedCrop(224),
                                     transforms.RandomHorizontalFlip(),
                                     transforms.ToTensor(),
                                     transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]),
        "val": transforms.Compose([transforms.Resize(256),
                                   transforms.CenterCrop(224),
                                   transforms.ToTensor(),
                                   transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])}

    train_dataset = Flowers(dataset_path=arg.dataset_path,
                            transforms=data_transform['train'])
    val_dataset = Flowers(dataset_path=arg.val_dataset_path,
                          transforms=data_transform['val'])

    train_dataloader = DataLoader(dataset=train_dataset,
                                  batch_size=arg.batch_size,
                                  shuffle=True,
                                  pin_memory=True,
                                  num_workers=arg.num_work,
                                  collate_fn=train_dataset.collate_fn)
    val_dataloader = DataLoader(dataset=val_dataset,
                                batch_size=arg.batch_size,
                                shuffle=False,
                                pin_memory=True,
                                num_workers=arg.num_work,
                                collate_fn=train_dataset.collate_fn)

    dataloaders = {'train': train_dataloader, 'val': val_dataloader}

    class_num = len(train_dataset.class2label)
    # 建立模型
    model = Vit(img_size=[224, 224],
                patch_size=16,
                embed_dim=192,
                depth=6,
                num_heads=4,
                num_classes=class_num)
    runner = Runner(arg, model, device)
    runner.run(dataloaders)


if __name__ == '__main__':
    arg = argparse.ArgumentParser()
    arg.add_argument('--work_dir', type=str, default='work_dir')
    arg.add_argument('--dataset_path', type=str, default='datasets/flower_data/train')
    arg.add_argument('--val_dataset_path', type=str, default='datasets/flower_data/val')
    arg.add_argument('--batch_size', type=int, default=16)
    arg.add_argument('--num_work', type=int, default=6)
    arg.add_argument('--work_flow', type=dict, default={'val': 1,'train': 3})
    arg.add_argument('--load_from', default='')
    arg.add_argument('--epochs', type=int, default=500)
    arg.add_argument('--save_epoch', type=int, default=2)
    arg.add_argument('--lr', type=float, default=0.001)
    arg.add_argument('--lrf', type=float, default=0.01)
    arg.add_argument('--device', default='cuda:0', help='device id (i.e. 0 or 0,1 or cpu)')
    opt = arg.parse_args()

    main(opt)
